// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ************************************************************************* //
//                          avtWellBoreFilter.C                              //
// ************************************************************************* //

#include <avtWellBoreFilter.h>

#include <vtkCellArray.h>
#include <vtkCharArray.h>
#include <vtkDataSet.h>
#include <vtkFieldData.h>
#include <vtkFloatArray.h>
#include <vtkMath.h>
#include <vtkPointData.h>
#include <vtkPolyData.h>
#include <vtkRectilinearGrid.h>
#include <vtkStructuredGrid.h>

#include <DebugStream.h>

#include <string>
#include <vector>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

static int cylinder_quality_levels[4] = {
    3,
    5,
    9,
    17
};

// ****************************************************************************
//  Method: avtWellBoreFilter constructor
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Wed Aug 27 14:59:19 PST 2008
//
// ****************************************************************************

avtWellBoreFilter::avtWellBoreFilter(WellBoreAttributes &wb_atts)
{
    atts = wb_atts;

    CalculateCylPts();
}


// ****************************************************************************
//  Method: avtWellBoreFilter destructor
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Wed Aug 27 14:59:19 PST 2008
//
// ****************************************************************************

avtWellBoreFilter::~avtWellBoreFilter()
{
    for (int detail=0; detail<MAX_DETAIL_LEVELS; detail++)
        delete [] cyl_pts[detail];

    return;
}


// ****************************************************************************
//  Method: avtWellBoreFilter::SetAttributes
//
//  Purpose:
//      Sets the attributes for this filter.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Wed Aug 27 14:59:19 PST 2008
//
// ****************************************************************************

void
avtWellBoreFilter::SetAttributes(const WellBoreAttributes &wb_atts)
{
    atts = wb_atts;
}


// ****************************************************************************
//  Method: avtWellBoreFilter::ExecuteData
//
//  Purpose:
//      Does the actual VTK code to modify the dataset.
//
//  Arguments:
//      inDR      The input data representation.
//
//  Returns:      The output data representation.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Wed Aug 27 14:59:19 PST 2008
//
//  Modifications:
//      Eric Brugger, Mon Nov 10 13:09:28 PST 2008
//      Added the ability to display well bore names and stems. 
//
//      Eric Brugger, Tue Aug 19 14:16:17 PDT 2014
//      Modified the class to work with avtDataRepresentation.
//
// ****************************************************************************

avtDataTree_p
avtWellBoreFilter::ExecuteDataTree(avtDataRepresentation *inDR)
{
    //
    // Get the VTK data set, the domain number, and the label.
    //
    vtkDataSet *inDS = inDR->GetDataVTK();
    int domain = inDR->GetDomain();
    std::string label = inDR->GetLabel();

    if (inDS->GetDataObjectType() != VTK_RECTILINEAR_GRID &&
        inDS->GetDataObjectType() != VTK_STRUCTURED_GRID)
    {
        EXCEPTION1(ImproperUseException,
                   "Expecting a rectilinear or curvilinear grid");
    }

    int dims[3];
    float *xpts = NULL, *ypts = NULL, *zpts = NULL, *pts = NULL;
    if (inDS->GetDataObjectType() == VTK_RECTILINEAR_GRID)
    {
        vtkRectilinearGrid *grid = vtkRectilinearGrid::SafeDownCast(inDS);
        if (grid->GetDataDimension() != 3)
        {
            EXCEPTION1(ImproperUseException, "Expecting a 3D grid");
        }
        xpts = vtkFloatArray::SafeDownCast(grid->GetXCoordinates())
            ->GetPointer(0);
        ypts = vtkFloatArray::SafeDownCast(grid->GetYCoordinates())
            ->GetPointer(0);
        zpts = vtkFloatArray::SafeDownCast(grid->GetZCoordinates())
            ->GetPointer(0);
        grid->GetDimensions(dims);
    }
    else if (inDS->GetDataObjectType() == VTK_STRUCTURED_GRID)
    {
        vtkStructuredGrid *grid = vtkStructuredGrid::SafeDownCast(inDS);
        if (grid->GetDataDimension() != 3)
        {
            EXCEPTION1(ImproperUseException, "Expecting a 3D grid");
        }
        pts = vtkFloatArray::SafeDownCast(grid->GetPoints()->GetData())
            ->GetPointer(0);
        grid->GetDimensions(dims);
    }

    int nWellBores = atts.GetNWellBores();
    const std::vector<int> wellBores = atts.GetWellBores();
    const std::vector<std::string> wellNames = atts.GetWellNames();

    if (nWellBores <= 0)
    {
        debug3 << "No well bores specified!" << endl;
        GetOutput()->GetInfo().GetValidity().InvalidateOperation();
        return NULL;
    }

    //
    // Determine the base index for the mesh.
    //
    vtkDataArray *arr = inDS->GetFieldData()->GetArray("base_index");
    int baseIndex[3] = {0, 0, 0};
    if (arr != NULL)
    {
        vtkIntArray *ar2 = (vtkIntArray *) arr;
        baseIndex[0] = ar2->GetValue(0);
        baseIndex[1] = ar2->GetValue(1);
        baseIndex[2] = ar2->GetValue(2);
    }

    //
    // Create the data sets and labels for each of the well bores.
    //
    vtkDataSet **out_ds = new vtkDataSet*[nWellBores];
    std::vector<std::string> labels;

    int iWellBore = 0;
    for (int i = 0; i < nWellBores; i++)
    {
        out_ds[i] = CreateWell(wellBores, iWellBore, wellNames[i],
                               xpts, ypts, zpts, pts, baseIndex, dims);

        //
        // Create the label.
        //
        char label[80];
        snprintf(label, 80, "%d", i);
        labels.push_back(label);
    }

    //
    // Turn the data sets into a data tree.
    //
    avtDataTree_p outDT = NULL;
    outDT = new avtDataTree(nWellBores, out_ds, domain, labels);

    for (int i = 0; i < nWellBores; i++)
    {
        out_ds[i]->Delete();
    }
    delete [] out_ds;

    //
    // Set the levels for the level mapper.
    //
    GetOutput()->GetInfo().GetAttributes().SetLabels(labels);

    return outDT;
}


// ****************************************************************************
//  Method: avtWellBoreFilter::RefashionDataObjectInfo
//
//  Purpose:
//      Allows the filter to change its output's data object information, which
//      is a description of the data object.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Wed Aug 27 14:59:19 PST 2008
//
// ****************************************************************************

void
avtWellBoreFilter::RefashionDataObjectInfo(void)
{
    // IF YOU SEE FUNNY THINGS WITH EXTENTS, ETC, YOU CAN CHANGE THAT HERE.
}


// ****************************************************************************
//  Method:  avtWellBoreFilter::CalculateCylPts
//
//  Purpose:
//    Precalculate points for cylinder geometry.
//
//  Programmer: Eric Brugger
//  Creation:   October 1, 2008
//
//  Modifications:
//
// ****************************************************************************

void
avtWellBoreFilter::CalculateCylPts(void)
{
    for (int detail=0; detail<MAX_DETAIL_LEVELS; detail++)
    {
        int cdetail = cylinder_quality_levels[detail];
        cyl_pts[detail] = new float[(cdetail+1)*4];

        for (int b=0; b<=cdetail; b++)
        {
            float theta = 2*M_PI * float(b) / float(cdetail);

            float dx = cos(theta);
            float dy = sin(theta);
            float dz = 0;

            cyl_pts[detail][b*4+0] = dx;
            cyl_pts[detail][b*4+1] = dy;
            cyl_pts[detail][b*4+2] = dz;
            cyl_pts[detail][b*4+3] = 0;
        }
    }
}

// ****************************************************************************
//  Method:  avtWellBoreFilter::GetPoint
//
//  Purpose:
//    Get the next point from the well bores.
//
//  Returns:    True if a point was present, false otherwise.
//
//  Arguments:
//    p         The next point.
//    wellBores The vector of well bores.
//    iWellBore The index of the next well bore.
//
//  Programmer: Eric Brugger
//  Creation:   October 1, 2008
//
//  Modifications:
//
// ****************************************************************************

bool
avtWellBoreFilter::GetPoint(int p[3], const std::vector<int> &wellBores,
    int &iWellBore)
{
    int i;
    for (i = 0; i < 3 && (p[i] = wellBores[iWellBore++]) != -1; i++)
        /* do nothing */;
    return (i == 3);
}

// ****************************************************************************
//  Method:  avtWellBoreFilter::CreateWell
//
//  Purpose:
//    Create the next well from the well description.
//
//  Returns:    A vtk poly data representation of the well.
//
//  Arguments:
//   wellBores  The vector of well bores.
//   iWellBore  The index of the next well bore.
//   wellName   The name of the well bore.
//   xpts       The x points of a rectilinear mesh.
//   ypts       The y points of a rectilinear mesh.
//   zpts       The z points of a rectilinear mesh.
//   pts        The points of a structured mesh.
//   baseIndex  The base index of the mesh.
//   dims       The dimensions of the mesh.
//
//  Programmer: Eric Brugger
//  Creation:   October 1, 2008
//
//  Modifications:
//      Eric Brugger, Mon Nov 10 13:09:28 PST 2008
//      I added logic to display a well stem. I also moved some code out of
//      this routine into the new routine AddWellSegment.  I added logic to
//      add a well name to the data set as character field data if we were
//      adding a well stem.
//
// ****************************************************************************

vtkPolyData *
avtWellBoreFilter::CreateWell(const std::vector<int> &wellBores,
    int &iWellBore, const std::string &wellName, float *xpts, float *ypts,
    float *zpts, float *pts, int baseIndex[3], int dims[3])
{
    WellBoreAttributes::WellRenderingMode renderingMode =
        atts.GetDrawWellsAs();

    vtkPolyData *polyData = vtkPolyData::New();

    vtkPoints *points = vtkPoints::New();
    polyData->SetPoints(points);
    points->Delete();

    vtkCellArray *lines = NULL;
    vtkCellArray *polys = NULL;
    vtkFloatArray *normals = NULL;
    if (renderingMode == WellBoreAttributes::Lines)
    {
        lines = vtkCellArray::New();
        polyData->SetLines(lines);
        lines->Delete();
    }
    else
    {
        polys = vtkCellArray::New();
        polyData->SetPolys(polys);
        polys->Delete();
        normals = vtkFloatArray::New();
        normals->SetName("Normals");
        normals->SetNumberOfComponents(3);
        polyData->GetPointData()->SetNormals(normals);
        normals->Delete();
    }

    int nx = dims[0];
    int nxy = dims[0] * dims[1];

    int nSegments = 0;
    int p1[3], p2[3];

    bool havePoints = GetPoint(p1, wellBores, iWellBore);
    if (havePoints)
    {
        //
        // If the first point is in this block, then add the well stem
        // and well name.
        //
        if (p1[0] >= baseIndex[0] && p1[0] < baseIndex[0] + dims[0] &&
            p1[1] >= baseIndex[1] && p1[1] < baseIndex[1] + dims[1] &&
            p1[2] >= baseIndex[2] && p1[2] < baseIndex[2] + dims[2])
        {
            WellBoreAttributes::WellAnnotation wellAnnotation =
                atts.GetWellAnnotation();

            //
            // Add the well stem.
            //
            if (wellAnnotation == WellBoreAttributes::StemOnly ||
                wellAnnotation == WellBoreAttributes::StemAndName)
            {
                float wellStemHeight = atts.GetWellStemHeight();

                //
                // Insert a point for the label.
                //
                double coord1[3], coord2[3];
                if (pts != NULL)
                {
                    int ndx = ((p1[2] - baseIndex[2]) * nxy +
                               (p1[1] - baseIndex[1]) * nx +
                               (p1[0] - baseIndex[0])) * 3;
                    coord1[0] = pts[ndx++];
                    coord1[1] = pts[ndx++];
                    coord1[2] = pts[ndx];
                }
                else
                {
                    coord1[0] = xpts[p1[0] - baseIndex[0]];
                    coord1[1] = ypts[p1[1] - baseIndex[1]];
                    coord1[2] = zpts[p1[2] - baseIndex[2]];
                }
                coord2[0] = coord1[0];
                coord2[1] = coord1[1];
                coord2[2] = coord1[2] - wellStemHeight;

                AddWellSegment(points, normals, lines, polys, nSegments,
                               coord2, coord1);
            }

            //
            // Add the well name.
            //
            if (wellAnnotation == WellBoreAttributes::NameOnly ||
                wellAnnotation == WellBoreAttributes::StemAndName)
            {
                vtkCharArray *arr = vtkCharArray::New();
                arr->SetNumberOfValues(wellName.length()+1);
                for (vtkIdType i = 0; i < (vtkIdType)wellName.length()+1; i++)
                {
                    arr->SetValue(i, wellName.c_str()[i]);
                }
                arr->SetName("well_name");
                polyData->GetFieldData()->AddArray(arr);
                arr->Delete();
            }
        }
        havePoints = GetPoint(p2, wellBores, iWellBore);
    }

    while (havePoints)
    {
        int nChange = 0;
        int nRemainingPoints = 0;
        int iRemainingPoints = 0;
        int i;
        for (i = 0; i < 3; i++)
        {
            int diff = p2[i] - p1[i];
            nChange += (diff != 0) ? 1 : 0;
            if (abs(diff) > abs(nRemainingPoints))
            {
                nRemainingPoints = diff;
                iRemainingPoints = i;
            }
        }
        if (nChange > 1 && abs(nRemainingPoints) > 1)
        {
            p1[0] = p2[0]; p1[1] = p2[1]; p1[2] = p2[2];
            havePoints = GetPoint(p2, wellBores, iWellBore);
        }
        else
        {
            if (p1[0] >= baseIndex[0] && p1[0] < baseIndex[0] + dims[0] &&
                p1[1] >= baseIndex[1] && p1[1] < baseIndex[1] + dims[1] &&
                p1[2] >= baseIndex[2] && p1[2] < baseIndex[2] + dims[2])
            {
                if (nRemainingPoints > 1)
                {
                    p2[iRemainingPoints] = p1[iRemainingPoints] + 1;
                }
                else if (nRemainingPoints < -1)
                {
                    p2[iRemainingPoints] = p1[iRemainingPoints] - 1;
                }
                if (p2[0] >= baseIndex[0] && p2[0] < baseIndex[0] + dims[0] &&
                    p2[1] >= baseIndex[1] && p2[1] < baseIndex[1] + dims[1] &&
                    p2[2] >= baseIndex[2] && p2[2] < baseIndex[2] + dims[2])
                {
                    double coord1[3], coord2[3];
                    if (pts != NULL)
                    {
                        int ndx = ((p1[2] - baseIndex[2]) * nxy +
                                   (p1[1] - baseIndex[1]) * nx +
                                   (p1[0] - baseIndex[0])) * 3;
                        coord1[0] = pts[ndx++];
                        coord1[1] = pts[ndx++];
                        coord1[2] = pts[ndx];
                        ndx = ((p2[2] - baseIndex[2]) * nxy +
                               (p2[1] - baseIndex[1]) * nx +
                               (p2[0] - baseIndex[0])) * 3;
                        coord2[0] = pts[ndx++];
                        coord2[1] = pts[ndx++];
                        coord2[2] = pts[ndx];
                    }
                    else
                    {
                        coord1[0] = xpts[p1[0] - baseIndex[0]];
                        coord1[1] = ypts[p1[1] - baseIndex[1]];
                        coord1[2] = zpts[p1[2] - baseIndex[2]];
                        coord2[0] = xpts[p2[0] - baseIndex[0]];
                        coord2[1] = ypts[p2[1] - baseIndex[1]];
                        coord2[2] = zpts[p2[2] - baseIndex[2]];
                    }

                    AddWellSegment(points, normals, lines, polys, nSegments,
                                   coord1, coord2);
                }
            }
        }

        if (nRemainingPoints > 1)
        {
            p1[iRemainingPoints]++;
            p2[iRemainingPoints] = p1[iRemainingPoints] + nRemainingPoints - 1;
        }
        else if (nRemainingPoints < -1)
        {
            p1[iRemainingPoints]--;
            p2[iRemainingPoints] = p1[iRemainingPoints] + nRemainingPoints + 1;
        }
        else
        {
            p1[0] = p2[0]; p1[1] = p2[1]; p1[2] = p2[2];
            havePoints = GetPoint(p2, wellBores, iWellBore);
        }
    }

    return polyData;
}

// ****************************************************************************
//  Method:  avtWellBoreFilter::AddWellSegment
//
//  Purpose:
//    Add a well segment to the current well.
//
//  Arguments:
//   points     The vtkPoints object to add points to.
//   normals    The vtkFloatArray object to add normals to. 
//   lines      The vtkCellArray object to add lines to. 
//   polys      The vtkCellArray object to add polygons to. 
//   nSegments  The number of segments added to the current well.
//   coord1     The first coordinate of the well segment.
//   coord2     The second coordinate of the well segment.
//
//  Programmer: Eric Brugger
//  Creation:   November 11, 2008
//
//  Modifications:
//
// ****************************************************************************

void
avtWellBoreFilter::AddWellSegment(vtkPoints *points, vtkFloatArray *normals,
    vtkCellArray *lines, vtkCellArray *polys, int &nSegments,
    double coord1[3], double coord2[3])
{
    WellBoreAttributes::WellRenderingMode renderingMode =
        atts.GetDrawWellsAs();
    WellBoreAttributes::DetailLevel cylinderQuality =
        atts.GetWellCylinderQuality();
    float cylinderRadius = atts.GetWellRadius();

    if (renderingMode == WellBoreAttributes::Lines)
    {
        points->InsertNextPoint(coord1);
        points->InsertNextPoint(coord2);

        vtkIdType line[2];
        line[0] = nSegments * 2;
        line[1] = nSegments * 2 + 1;
        lines->InsertNextCell(2, line);

        nSegments++;
    }
    else
    {
        float vc[3] = {static_cast<float>(coord2[0] - coord1[0]),
                       static_cast<float>(coord2[1] - coord1[1]),
                       static_cast<float>(coord2[2] - coord1[2])};
        float va[3];
        float vb[3];

        float vc_len = vtkMath::Normalize(vc);
        if (vc_len != 0)
        {
            vtkMath::Perpendiculars(vc, va,vb, 0);

            int cdetail =
                cylinder_quality_levels[cylinderQuality];
            for (int b = 0; b < cdetail; b++)
            {
                float *u = &(cyl_pts[cylinderQuality][b*4]);
                float v[3];

                v[0] = va[0]*u[0] + vb[0]*u[1];
                v[1] = va[1]*u[0] + vb[1]*u[1];
                v[2] = va[2]*u[0] + vb[2]*u[1];

                float normal[3];
                normal[0] = v[0];
                normal[1] = v[1];
                normal[2] = v[2];
                normals->InsertNextTuple(normal);
                normals->InsertNextTuple(normal);

                double coord[3];
                coord[0] = coord1[0] + cylinderRadius * v[0];
                coord[1] = coord1[1] + cylinderRadius * v[1];
                coord[2] = coord1[2] + cylinderRadius * v[2];
                points->InsertNextPoint(coord);
                coord[0] = coord2[0] + cylinderRadius * v[0];
                coord[1] = coord2[1] + cylinderRadius * v[1];
                coord[2] = coord2[2] + cylinderRadius * v[2];
                points->InsertNextPoint(coord);
            }

            vtkIdType poly[4];
            for (int b = 0; b < cdetail; b++)
            {
                int offset = nSegments * cdetail * 2;
                poly[0] = offset + b * 2;
                poly[1] = offset + ((b + 1) % cdetail) * 2;
                poly[2] = offset + ((b + 1) % cdetail) * 2 + 1;
                poly[3] = offset + b * 2 + 1;
                polys->InsertNextCell(4, poly);
            }

            nSegments++;
        }
    }
}
